"""外部认证系统集成模块。

本模块处理与外部认证系统（如Auth0、Firebase Auth、Keycloak等）的集成，
包括令牌验证和用户信息获取等功能。
"""

import logging
from typing import Optional, Dict, Any
import httpx
from jose import jwt, JWTError

from app.core.config import settings
from app.core.logging import logger

class ExternalAuthClient:
    """外部认证系统客户端。"""
    
    def __init__(self):
        """初始化外部认证客户端。"""
        self.jwks_client = None
        self.issuer = settings.EXTERNAL_AUTH_ISSUER if hasattr(settings, 'EXTERNAL_AUTH_ISSUER') else None
        self.audience = settings.EXTERNAL_AUTH_AUDIENCE if hasattr(settings, 'EXTERNAL_AUTH_AUDIENCE') else None
        
        # 如果配置了JWKS URI，则初始化JWKS客户端
        if hasattr(settings, 'EXTERNAL_AUTH_JWKS_URI') and settings.EXTERNAL_AUTH_JWKS_URI:
            self.jwks_client = httpx.Client()
            
        logger.info("external_auth_client_initialized")
    
    async def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
        """验证外部认证系统的JWT令牌。

        Args:
            token: JWT令牌字符串

        Returns:
            Optional[Dict[str, Any]]: 如果令牌有效则返回载荷，否则返回None
        """
        # 检查外部认证是否启用
        if not settings.EXTERNAL_AUTH_ENABLED:
            logger.debug("external_auth_disabled")
            return None
            
        if not token:
            logger.warning("external_auth_token_missing")
            return None
            
        try:
            # 如果配置了JWKS URI，则使用JWKS验证
            if self.jwks_client and hasattr(settings, 'EXTERNAL_AUTH_JWKS_URI'):
                return await self._verify_with_jwks(token)
            else:
                # 否则使用密钥验证
                return await self._verify_with_key(token)
        except Exception as e:
            logger.error("external_auth_token_verification_failed", error=str(e))
            return None
    
    async def _verify_with_jwks(self, token: str) -> Optional[Dict[str, Any]]:
        """使用 JWKS 动态密钥验证令牌。

        Args:
            token (str): 要验证的 JWT 令牌

        Returns:
            Optional[Dict[str, Any]]: 验证成功返回解析后的载荷，失败返回 None
        """
        try:
            # 获取 JWKS 密钥集
            jwks_response = self.jwks_client.get(settings.EXTERNAL_AUTH_JWKS_URI)
            jwks = jwks_response.json()
            
            # 解码头部以获取 kid
            header = jwt.get_unverified_header(token)
            kid = header.get("kid")
            
            if not kid:
                logger.warning("external_auth_token_missing_kid")
                return None
                
            # 从 JWKS 密钥集中查找匹配的密钥
            rsa_key = {}
            for key in jwks["keys"]:
                if key["kid"] == kid:
                    rsa_key = {
                        "kty": key["kty"],
                        "kid": key["kid"],
                        "use": key["use"],
                        "n": key["n"],
                        "e": key["e"]
                    }
                    break
                    
            if not rsa_key:
                logger.warning("external_auth_jwk_not_found", kid=kid)
                return None
                
            # 使用找到的 RSA 密钥验证令牌
            payload = jwt.decode(
                token,
                rsa_key,
                algorithms=["RS256"],
                audience=self.audience,
                issuer=self.issuer
            )
            
            return payload
        except JWTError as e:
            logger.error("external_auth_jwks_verification_failed", error=str(e))
            return None

    async def _verify_with_key(self, token: str) -> Optional[Dict[str, Any]]:
        """使用共享密钥验证令牌。

        Args:
            token (str): 要验证的 JWT 令牌

        Returns:
            Optional[Dict[str, Any]]: 验证成功返回解析后的载荷，失败返回 None
        """
        try:
            # 使用共享密钥验证令牌
            payload = jwt.decode(
                token,
                settings.JWT_SECRET_KEY,  # 外部认证系统共享密钥
                algorithms=[settings.JWT_ALGORITHM],
                audience=self.audience,
                issuer=self.issuer
            )
            
            return payload
        except JWTError as e:
            logger.error("external_auth_key_verification_failed", error=str(e))
            return None

# 创建单例实例
external_auth_client = ExternalAuthClient()